#include <functional>
#include <vector>
#include <tuple>
#include <algorithm>

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
public:
    void recoverTree(TreeNode* root) {
        TreeNode* last = nullptr;
        std::vector<std::tuple<TreeNode*,TreeNode*>> invalid;
        std::function<void(TreeNode*)> traverse = [&](TreeNode* root) mutable {
            if (!root) return;
            traverse(root->left);
            if (last && last->val>root->val)
                invalid.emplace_back(last, root);
            last = root;
            traverse(root->right);
        };
        traverse(root);
        if (invalid.size()==2){
            std::swap(std::get<0>(invalid[0])->val, std::get<1>(invalid[1])->val);
        } else if (invalid.size() == 1){
            auto little = std::get<0>(invalid[0])->val;
            TreeNode* lastLittleNode = nullptr;
            auto large = std::get<1>(invalid[0])->val;
            TreeNode* firstLargeNode = nullptr;
            std::function<void(TreeNode*)> traverse = [&](TreeNode* root) mutable{
                if (!root) return;
                if (root->val == little) lastLittleNode = root;
                if (root->val == large && !firstLargeNode) firstLargeNode = root; 
                traverse(root->left);
                traverse(root->right);
            };
            traverse(root);
            std::swap(firstLargeNode->val, lastLittleNode->val);
        }
    }
};